"""PDF处理相关节点"""

from pocketflow import Node
from ..utils.parse_pdf import parse_pdf
from ..utils.chinese_chunk_utils import split_text_into_chunks

class ParsePDFNode(Node):
    """解析PDF文件的节点"""
    
    def prep(self, shared):
        return shared["pdf_path"]

    def exec(self, pdf_path):
        return parse_pdf(pdf_path)

    def post(self, shared, prep_res, exec_res):
        shared["pdf_text"] = exec_res
        return "default"

class ChunkTextNode(Node):
    """将长文本智能分割为块，特别优化中文段落处理
    
    改进特性：
    1. 识别多层级标题结构（章节、子章节、列表项）
    2. 保持列表项与其标题的完整性
    3. 在句子边界处分割，避免截断语义
    4. 适应中文文档的较大chunk_size（1200字符）
    """
    
    def __init__(self, max_chunk_size: int = 1200, **kwargs):
        """
        Args:
            max_chunk_size: 最大块大小（字符数），默认1200适合中文段落
        """
        super().__init__(**kwargs)
        self.max_chunk_size = max_chunk_size
    
    def prep(self, shared):
        pdf_text = shared["pdf_text"]
        print(f"开始智能分块，文本总长度: {len(pdf_text)} 字符")
        return pdf_text
    
    def exec(self, pdf_text):
        # 使用新的中文智能分块算法
        chunks = split_text_into_chunks(pdf_text, max_chunk_size=self.max_chunk_size)
        print(f"分块完成，共 {len(chunks)} 个块")
        
        # 打印每个块的统计信息
        for chunk in chunks:
            print(f"  块{chunk['chunk_id']}: {chunk['title'][:30]}... ({chunk['char_count']}字符)")
        
        return chunks
    
    def post(self, shared, prep_res, exec_res):
        shared["text_chunks"] = exec_res
        shared["chunk_checkpoints"] = {}  # 初始化块检查点存储
        
        # 保存分块统计信息
        shared["chunk_stats"] = {
            "total_chunks": len(exec_res),
            "total_chars": len(prep_res),
            "avg_chunk_size": len(prep_res) // len(exec_res) if exec_res else 0,
            "chunks_detail": [
                {
                    "id": chunk["chunk_id"],
                    "title": chunk["title"],
                    "size": chunk["char_count"]
                }
                for chunk in exec_res
            ]
        }
        
        return "default"
